import requests
import json
import argparse
from rich.console import Console
from rich.table import Table

# 初始化控制台
console = Console()

def query_account(username, password, api_url="http://127.0.0.1:8000/api/account"):
    """
    查询TQSDK账户信息
    
    Args:
        username: TQSDK用户名
        password: TQSDK密码
        api_url: API地址
    
    Returns:
        查询结果
    """
    payload = {
        "username": username,
        "password": password
    }
    headers = {
        "Content-Type": "application/json"
    }
    
    try:
        console.print(f"正在查询账户信息...", style="bold blue")
        try:
            response = requests.post(api_url, data=json.dumps(payload), headers=headers, timeout=10)
            
            if response.status_code != 200:
                console.print(f"查询失败: HTTP状态码 {response.status_code}", style="bold red")
                if response.status_code == 404:
                    console.print(f"API地址不存在，请检查服务器是否启动或URL是否正确", style="yellow")
                elif response.status_code == 401:
                    console.print(f"认证失败，请检查用户名和密码", style="yellow")
                elif response.status_code == 500:
                    console.print(f"服务器内部错误，请查看服务器日志", style="yellow")
                return None
            
            result = response.json()
            return result
        except requests.exceptions.ConnectionError:
            console.print(f"连接失败: 无法连接到API服务器 {api_url}", style="bold red")
            console.print(f"请确保API服务器已启动且地址正确", style="yellow")
            return None
        except requests.exceptions.Timeout:
            console.print(f"连接超时: API服务器响应时间过长", style="bold red")
            return None
    except Exception as e:
        console.print(f"查询出错: {str(e)}", style="bold red")
        return None

def display_account_info(result):
    """
    显示账户信息
    
    Args:
        result: API返回的结果
    """
    if not result or not result.get("success"):
        console.print(f"查询失败: {result.get('message', '未知错误')}", style="bold red")
        return
    
    # 显示账户资金信息
    account = result["data"]["account"]
    account_table = Table(title="账户资金信息", show_lines=True)
    account_table.add_column("项目", style="cyan")
    account_table.add_column("数值", style="green")
    
    account_table.add_row("可用资金", f"{account['available']:,.2f}")
    account_table.add_row("账户权益", f"{account['balance']:,.2f}")
    account_table.add_row("币种", account['currency'])
    account_table.add_row("占用保证金", f"{account['margin']:,.2f}")
    account_table.add_row("风险度", f"{account['risk_ratio']:.4f}")
    account_table.add_row("昨日账户权益", f"{account['pre_balance']:,.2f}")
    account_table.add_row("浮动盈亏", f"{account['float_profit']:,.2f}")
    account_table.add_row("平仓盈亏", f"{account['close_profit']:,.2f}")
    
    console.print(account_table)
    
    # 显示持仓信息
    positions = result["data"]["positions"]
    if positions:
        position_table = Table(title="持仓信息", show_lines=True)
        position_table.add_column("合约", style="cyan")
        position_table.add_column("多头持仓", style="green")
        position_table.add_column("空头持仓", style="red")
        position_table.add_column("今日多开", style="green")
        position_table.add_column("今日空开", style="red")
        position_table.add_column("多头盈亏", style="green")
        position_table.add_column("空头盈亏", style="red")
        position_table.add_column("多头均价", style="green")
        position_table.add_column("空头均价", style="red")
        
        for position in positions:
            position_table.add_row(
                position["symbol"],
                str(position["pos_long"]),
                str(position["pos_short"]),
                str(position["volume_long_today"]),
                str(position["volume_short_today"]),
                f"{position['float_profit_long']:,.2f}" if position['float_profit_long'] is not None else "0.00",
                f"{position['float_profit_short']:,.2f}" if position['float_profit_short'] is not None else "0.00",
                f"{position['position_price_long']:,.2f}" if position['position_price_long'] is not None and position['position_price_long'] > 0 else "0.00",
                f"{position['position_price_short']:,.2f}" if position['position_price_short'] is not None and position['position_price_short'] > 0 else "0.00"
            )
        
        console.print(position_table)
    else:
        console.print("当前账户没有持仓", style="yellow")

def main():
    parser = argparse.ArgumentParser(description="TQSDK账户查询客户端")
    parser.add_argument("--username", "-u", help="TQSDK用户名")
    parser.add_argument("--password", "-p", help="TQSDK密码")
    parser.add_argument("--url", default="http://127.0.0.1:8000/api/account", help="API地址")
    parser.add_argument("--default", "-d", action="store_true", help="使用默认凭据")
    
    args = parser.parse_args()
    
    # 使用默认凭据或命令行参数
    if args.default:
        username = "cps168"
        password = "alibaba"
        console.print("使用默认凭据进行查询...", style="bold yellow")
    else:
        if not args.username or not args.password:
            parser.print_help()
            console.print("\n错误: 必须提供用户名和密码，或使用 --default 参数使用默认凭据", style="bold red")
            return
        username = args.username
        password = args.password
    
    result = query_account(username, password, args.url)
    if result:
        display_account_info(result)

def print_welcome():
    """打印欢迎信息和使用说明"""
    console.print("=" * 60, style="bold blue")
    console.print("TQSDK 账户查询客户端", style="bold cyan")
    console.print("=" * 60, style="bold blue")
    console.print("使用方法:")
    console.print("  1. 使用自定义凭据: python account_client.py -u 用户名 -p 密码")
    console.print("  2. 使用默认凭据: python account_client.py --default")
    console.print("  3. 指定API地址: python account_client.py -u 用户名 -p 密码 --url http://api地址")
    console.print("=" * 60, style="bold blue")
    console.print("")

if __name__ == "__main__":
    print_welcome()
    main()